# 1. 导入必要库
import numpy as np
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier, plot_tree, export_text
import matplotlib.pyplot as plt

plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']  # 指定默认字体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# 2. 加载并准备数据
iris = load_iris()
X = iris.data  # 特征矩阵（萼片长宽、花瓣长宽）
y = iris.target  # 目标标签（3种鸢尾花）

# 3. 划分训练集和测试集（80%训练，20%测试）
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

# 4. 创建ID3决策树模型（关键参数设置）
id3_tree = DecisionTreeClassifier(
    criterion="entropy",  # 使用信息熵（ID3核心）
    max_depth=3,          # 限制树深度防过拟合
    random_state=42       # 固定随机性确保结果可复现
)

# 5. 训练模型
id3_tree.fit(X_train, y_train)

# 6. 评估模型性能
train_accuracy = id3_tree.score(X_train, y_train)
test_accuracy = id3_tree.score(X_test, y_test)
print(f"训练集准确率: {train_accuracy:.2f}")
print(f"测试集准确率: {test_accuracy:.2f}")

# 7. 可视化决策树
plt.figure(figsize=(15, 10))
plot_tree(
    id3_tree,
    feature_names=iris.feature_names,  # 显示特征名
    class_names=iris.target_names,     # 显示类别名
    filled=True,                       # 颜色填充区分类别
    rounded=True,                      # 圆角矩形节点
    impurity=True,                     # 显示节点不纯度（熵）
)
plt.title("ID3决策树可视化（基于信息熵）")
plt.show()

# 8. 文本形式输出决策规则（无需图形依赖）
tree_rules = export_text(
    id3_tree,
    feature_names=iris.feature_names
)
print("决策规则路径：\n", tree_rules)

# 9. 预测新样本
new_sample = [[5.0, 3.2, 1.5, 0.3]]  # 新数据点
predicted_class = id3_tree.predict(new_sample)
probabilities = id3_tree.predict_proba(new_sample)
print(f"预测类别: {iris.target_names[predicted_class[0]]}")
print(f"类别概率: {dict(zip(iris.target_names, probabilities[0]))}")